import pandas as pd
from sklearn.model_selection import train_test_split


# 按类别（标签）划分数据集
def split_dataset_by_labels(df):
    # 存储划分后的训练集和ReN-GAN训练集
    train_subset = pd.DataFrame()
    ren_gan_subset = pd.DataFrame()

    # 每类流量的标签列名
    label_columns = ['*label_Dos', '*label_Probe', '*label_R2L', '*label_U2R', '*label_normal']

    # 对每一类标签进行划分
    for label in label_columns:
        # 根据标签列选择出符合该标签的所有样本
        label_data = df[df[label] == 1]

        # 按50%的比例划分为训练集和ReN-GAN训练集
        train_data, ren_gan_data = train_test_split(label_data, test_size=0.5, random_state=42)

        # 将划分后的数据分别加入训练集和ReN-GAN训练集
        train_subset = pd.concat([train_subset, train_data], axis=0)
        ren_gan_subset = pd.concat([ren_gan_subset, ren_gan_data], axis=0)

    return train_subset, ren_gan_subset

# 加载数据集
file_path = './data/data_with_normal.csv'
df = pd.read_csv(file_path)

# 使用加载的数据集进行划分
train_subset, ren_gan_subset = split_dataset_by_labels(df)

# 输出划分结果
print(f"训练集大小: {train_subset.shape}")
print(f"ReN-GAN训练集大小: {ren_gan_subset.shape}")

# 保存训练集和ReN-GAN训练集
train_subset.to_csv("./data/train_subset.csv", index=False)
ren_gan_subset.to_csv("./data/ren_gan_subset.csv", index=False)
